import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.gridspec import GridSpec

plt.rcParams['font.sans-serif'] = ['Times New Roman', 'Heiti TC', 'Microsoft YaHei', 'SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

species = ("2025", "2030", "2035", "2040", "2045", "2050")

def read_csv(csv_path):
    df = pd.read_csv(csv_path, index_col=0)
    df.rename(index={"Hydro_Nonpumped": "Hydro", "Central_PV": "PV"}, inplace=True)
    df.columns = species
    energy = ["Coal", "Gas", "Nuclear", "Hydro", "PV", "Wind"]
    weights = {row: np.array(df.loc[row].values / 10) for row in energy}
    normalized_df = df.div(df.sum(axis=0), axis=1)
    normalized_weights = {row: normalized_df.loc[row].values for row in energy}
    total = df.sum(axis=0).values / 10  # 总装机量（亿千瓦）
    return weights, normalized_weights, total

width = 0.8
colors = ['#000000', '#ffdbba', '#f4a261', '#00c0ff', '#ffd800', '#39b371']

def plot_in_axes(ax, species, weights):
    bottom = np.zeros(len(species))
    for i, (boolean, weight) in enumerate(weights.items()):
        ax.bar(species, weight, width, label=boolean, bottom=bottom, color=colors[i])
        bottom += weight

# 设置布局：2行3列，但手动定义 GridSpec 实现第一行跨列
fig = plt.figure(figsize=(10, 5))
gs = GridSpec(2, 3, height_ratios=[1, 1])

ax_top = fig.add_subplot(gs[0, :])  # 第一行跨越所有列
ax_bottom = [fig.add_subplot(gs[1, i]) for i in range(3)]

# 读取数据
_, _, total1 = read_csv("data-1fix.csv")
_, _, total2 = read_csv("data-1v1g.csv")
_, _, total3 = read_csv("data-1v2g.csv")

# 绘制顶部折线图
ax_top.plot(species, total1, marker='o', label='BAU-LOW-UC', color="#bb5b4b")
ax_top.plot(species, total2, marker='s', label='BAU-LOW-V1G', color="#507f7f")
ax_top.plot(species, total3, marker='^', label='BAU-LOW-V2G', color="#1f3481")
ax_top.set_ylim(20, max(total1.max(), total2.max(), total3.max()) * 1.1)
ax_top.set_ylabel("Installed Capacity(MMW)")
# ax_top.set_title("不同场景下总装机量随时间变化")
ax_top.grid(axis='y', linestyle='-', alpha=0.6)
ax_top.set_xlabel("(a)")
# ax_top.legend(loc='upper left')

def plot_csv_in_col(csv_path, ax, xlabel):
    w, nw, _ = read_csv(csv_path)
    ax.set_ylim(top=1)
    ax.set_xlabel(xlabel)
    ax.grid(axis='y', linestyle='-', alpha=0.6)
    plot_in_axes(ax, species, nw)

# 底部三张图
plot_csv_in_col("data-1fix.csv", ax_bottom[0], "(b)")
plot_csv_in_col("data-1v1g.csv", ax_bottom[1], "(c)")
plot_csv_in_col("data-1v2g.csv", ax_bottom[2], "(d)")

ax_bottom[0].set_ylabel("Installed Capacity Share by Source")

# 添加统一图例到右侧
handles_bottom, labels_bottom = ax_bottom[0].get_legend_handles_labels()
handles_top, labels_top = ax_top.get_legend_handles_labels()

# 合并两个图例的handles和labels
all_handles = handles_bottom + handles_top
all_labels = labels_bottom + labels_top

# 添加统一的图例
fig.legend(all_handles, all_labels, loc='center left', bbox_to_anchor=(0.84, 0.5))

# 调整布局，预留右侧空间给图例
fig.tight_layout(rect=[0, 0, 0.85, 0.95])

plt.show()
